# coding: utf8

from PyPDF2 import PdfFileReader, PdfFileWriter


def read_info(input_pdf='demo.pdf'):
    """
    使用PyPDF2.PdfFileReader读取PDF文件
    获取文档的各类数据
    """

    # 获取 PdfFileReader 对象
    pdfdoc= PdfFileReader(input_pdf)
    # 或者这个方式：pdfFileReader = PdfFileReader(open(readFile, 'rb'))

    # 获取文档信息
    info = pdfdoc.getDocumentInfo()
    print('doc info = %s' % info)
    print(f'''
    作者: {info.author}
    创建者: {info.creator}
    制作者: {info.producer}
    主题: {info.subject}
    标题: {info.title}
    页数: {pdfdoc.getNumPages()}
    加密: {pdfdoc.isEncrypted}
    ''')

    # 获取页面布局
    pageLayout = pdfdoc.getPageLayout()
    print('pageLayout = %s ' % pageLayout)

    # 获取页模式
    pageMode = pdfdoc.getPageMode()
    print('pageMode = %s' % pageMode)
    xmpMetadata = pdfdoc.getXmpMetadata()
    print('xmpMetadata  = %s ' % xmpMetadata)

    # 获取页对象
    for index, page in enumerate(pdfdoc.pages):
        print('页：{} , page: {}'.format(index, page))
        # maybe result in chaos chars
        print("文档文本：", page.extractText())

    return pdfdoc


def split(input_pdf, page_group_tuple=None):
    """
    write each page of readpdf to a single pdf file: split_pdf_[pageno].pdf
    :param input_pdf: source pdf file
    :split_page_no_list: list of list of page number
    """
    # create reader to read, writers for adding pages to split pdf
    pdf_reader = PdfFileReader(input_pdf)
    pdf_writer_list = [PdfFileWriter() for _ in page_group_tuple]

    # add page to writer
    # max_page_num = pdf_reader.getNumPages()
    for j, page_list in enumerate(page_group_tuple):
        for pno in page_list:
            # if pno not in range(max_page_num):
            #     print('invalid page no: {}'.format(pno))
            #     continue
            # page = pdf_reader.getPage(pno)
            pdf_writer_list[j].addPage(pdf_reader.getPage(pno))

    # write to pdf file
    for fi, writer in enumerate(pdf_writer_list):
        split_file_name = 'split_pdf_{:02d}.pdf'.format(fi+1)
        with open(split_file_name, 'wb') as fp:
            writer.write(fp)


def merge(input_pdf_list, merge_pdf='merge_00.pdf'):
    # create readers, writer
    readers = [PdfFileReader(pdf) for pdf in input_pdf_list]
    merge_writer = PdfFileWriter()

    # Write reader.pages to merge_writer
    for reader in readers:
        for i in range(reader.getNumPages()):
            merge_writer.addPage(reader.getPage(i))

    # write to merger pdf
    with open(merge_pdf, 'wb') as fp:
        merge_writer.write(fp)


def add_watermark(input_pdf, watermark_pdf, out_pdf='temp.pdf'):
    # read watermark and input pdf
    wm_reader = PdfFileReader(watermark_pdf, strict=False)
    watermark = wm_reader.getPage(0)
    pdf_reader = PdfFileReader(input_pdf, strict=False)
    # add watermark to writer
    pdf_writer = PdfFileWriter()
    for page in pdf_reader.pages:
        # page = pdf_reader.getPage(pageno)
        page.mergePage(watermark)
        page.compressContentStreams()
        pdf_writer.addPage(page)
    # write to pdf file
    pdf_writer.write(open(out_pdf, 'wb'))


def encrypt(input_pdf, out_pdf='temp_encrypt.pdf', password='123456'):
    pdf_reader = PdfFileReader(input_pdf)
    pdf_writer = PdfFileWriter()
    # add pages to writer
    for pno in range(pdf_reader.getNumPages()):
        pdf_writer.addPage(pdf_reader.getPage(pno))
    # add password
    # pdf_writer.encrypt(user_pwd=password, owner_pwd=password, use_128bit=True)
    pdf_writer.encrypt(user_pwd=password, use_128bit=True)
    # write to file
    with open(out_pdf, 'wb') as fp:
        pdf_writer.write(fp)


def decrypt(input_pdf, out_pdf, password=None):
    # read and decrypt pdf
    pdf_reader = PdfFileReader(open(input_pdf, 'rb'))
    pdf_reader.decrypt(password)
    # add pages to writer
    pdf_writer = PdfFileWriter()
    for page in pdf_reader.pages:
        pdf_writer.addPage(page)
    # write to file
    with open(out_pdf, 'wb') as fp:
        pdf_writer.write(fp)


def rotation(input_pdf=None, writepdf='temp_rotation.pdf', rotation_dict=None):
    # create reader and writer
    pdf_reader = PdfFileReader(open(input_pdf, 'rb'))
    pdf_writer = PdfFileWriter()
    # rotate page and add page to writer
    for pno, page in enumerate(pdf_reader.pages):
        if pno in rotation_dict:
            pdf_writer.addPage(page.rotateClockwise(rotation_dict[pno]))
        else:
            pdf_writer.addPage(page)
    # write to file
    with open(writepdf, 'wb') as fp:
        pdf_writer.write(fp)


def remove_images(input_pdf, out_pdf='temp_without_images.pdf'):
    """由于识别问题，容易造成文字不过走样，如反转、漏字等"""
    # 读取原文档
    pdf_reader = PdfFileReader(open(input_pdf, 'rb'))
    # 添加到输出文档对象
    pdf_writer = PdfFileWriter()
    [pdf_writer.addPage(page) for page in pdf_reader.pages]
    # 去除图片
    pdf_writer.removeImages()
    # 将输出文档写入文件
    with open(out_pdf, 'wb') as fp:
        pdf_writer.write(fp)
    # rotation(out_pdf, rotation_dict={0: 180})
